

from envs import next_step_back_mp_coo_large_scale_twice_timeout_noprofiling, next_step_back_stable_mp_coo_large_scale_twice_timeout_noprofiling

import random
import numpy as np
import copy


def setup_seed(seed):
     np.random.seed(seed)
     random.seed(seed)
     
setup_seed(20230430)


# sparse_name_list = ['gen2', '176bit', 'bcsstk13', 'biplane-9', 
#                     'ch7-8-b3', 'luxembourg_osm', 'cond-mat-2003', 'astro-ph', 
#                     '3D_28984_Tetra', 'rajat25', 'c-69', 'rajat16', 
#                     'net4-1', 'tsyl201', 'matrix-new_3', 'ct20stif']

sparse_name_list = ['gen2', '176bit', 'bcsstk13', 'biplane-9', 't2d_q9', 'fv2', 'fv3', 'big_dual', 'ch7-8-b3', 'luxembourg_osm', 'cond-mat-2003', 'astro-ph', 'wing', 'TSOPF_RS_b39_c7', 'airfoil_2d', 'jan99jac120sc', '3D_28984_Tetra', 'rajat25', 'c-69', 'rajat16', 'Trec13', 'Zd_Jac2_db', 'bibd_17_8', 'image_interp', 'net4-1', 'tsyl201', 'matrix-new_3', 'ct20stif', 'stat96v2', 'belgium_osm', 'Chevron3', 'parabolic_fem']

rank = 0
world_size = 5

block_num, thread_num, element_num = 512, 256, 2

# sparse_name_list = ['ch7-8-b3']
# sparse_name = 'ch7-8-b3'

mean_gflops_list = []
std_gflops_list = []

for sparse_name in sparse_name_list:
    try:
        init_gflops = next_step_back_mp_coo_large_scale_twice_timeout_noprofiling(sparse_name, block_num, thread_num, element_num, rank, world_size)

    except Exception as e:
        
        flag_continue = True 
        while flag_continue:
            try:
                init_gflops = next_step_back_stable_mp_coo_large_scale_twice_timeout_noprofiling(sparse_name, block_num, thread_num, element_num, rank, world_size) # get next state

                flag_continue = False
            except:
                print('large: ', sparse_name, " ", block_num, " ", thread_num, " ", element_num)

    print('first: ', init_gflops)

    parameters_scale = np.array([(128, 4096, 128), (32, 1024, 32), (1, 32, 1)])

    gflops_list = []
    parameters_num = 32
    program_feature = [0, 0, 0]

    # for i in range(30):
        
        # for j in range(3):    
            # policy_select_value = random.randint(0, parameters_num-1)  
            # program_feature[j] = policy_select_value * parameters_scale[j][2] \
                                    # + parameters_scale[j][0] # get
    for i in range(128, 4098, 128 * 8):
        for j in range(32, 1026, 32 * 8):
            for k in range(1, 33, 1 * 8):
                    
                program_feature[0] = i
                program_feature[1] = j
                program_feature[2] = k
                    
                try:
                    gflops = next_step_back_mp_coo_large_scale_twice_timeout_noprofiling(sparse_name, program_feature[0], program_feature[1], 
                                                                                        program_feature[2], rank, world_size)
                except Exception as e:
                    
                    flag_continue = True 
                    while flag_continue:
                        try:
                            gflops = next_step_back_stable_mp_coo_large_scale_twice_timeout_noprofiling(sparse_name, program_feature[0], program_feature[1], 
                                                                                                            program_feature[2], rank, world_size) # get next state
                            flag_continue = False
                        except:
                            print('large: ', sparse_name, " ", program_feature[0], program_feature[1],  program_feature[2],)
                gflops_list.append([gflops, copy.deepcopy(program_feature)])
                print(f'{sparse_name} | config:', program_feature, '| gflops: ', gflops)


    # gflops_list = np.array(gflops_list)
    
    print(sparse_name)
    mid_max = max(gflops_list, key=lambda x: x[0])
    print(mid_max)
    

    program_feature[0] = mid_max[1][0]
    program_feature[1] = mid_max[1][1]
    
    final_list = []
    for k in range(1, 33):
        program_feature[2] = k
            
        try:
            gflops = next_step_back_mp_coo_large_scale_twice_timeout_noprofiling(sparse_name, program_feature[0], program_feature[1], 
                                                                                program_feature[2], rank, world_size)
        except Exception as e:
            
            flag_continue = True 
            while flag_continue:
                try:
                    gflops = next_step_back_stable_mp_coo_large_scale_twice_timeout_noprofiling(sparse_name, program_feature[0], program_feature[1], 
                                                                                                    program_feature[2], rank, world_size) # get next state
                    flag_continue = False
                except:
                    print('large: ', sparse_name, " ", program_feature[0], program_feature[1],  program_feature[2],)
        final_list.append([gflops, copy.deepcopy(program_feature)])
        print(f'{sparse_name} | config:' + str(program_feature) + ' | gflops: ' + str(gflops))
    
    final_max = max(final_list, key=lambda x: x[0])
    print(final_max)
    
    fo = open("record_best.txt", "a")
    fo.write(f'{sparse_name} | config:' +  str(final_max[1]) + ' | gflops: ' + str(final_max[0]) + '\n')
    fo.close() # 关闭文件
    
    print(f'{sparse_name} | config:', final_max[1], '| gflops: ', final_max[0])
    